package mcpTool

import (
	"buding-ai/internal/repository/model"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/mcp"
	"github.com/sashabaranov/go-openai"
	"strings"
)

type McpClientList []*McpClient

func (c McpClientList) AllTools() []openai.Tool {
	tools := make([]openai.Tool, 0)
	for _, mcpClient := range c {
		tools = append(mcpClient.tools, tools...)
	}
	return tools
}

func (c McpClientList) Close() {
	for _, mcpClient := range c {
		mcpClient.client.Close()
	}
}

func (c McpClientList) GetClient() {
	for _, mcpClient := range c {
		mcpClient.client.Close()
	}
}

func (c McpClientList) getClientWithName(name string) *McpClient {
	for _, val := range c {
		if val.name == name {
			return val
		}
	}
	return nil
}

func (c McpClientList) Call(name, param string) (string, error) {
	split := strings.Split(name, "-")
	mcpClient := c.getClientWithName(split[0])
	if mcpClient == nil {
		return "", errors.New(`tool client not found`)
	}

	toolRequest := mcp.CallToolRequest{
		Request: mcp.Request{
			Method: "tools/call",
		},
	}
	toolRequest.Params.Name = split[1]
	params := make(map[string]any)
	err := json.Unmarshal([]byte(param), &params)
	if err != nil {
		return "", errors.New(fmt.Sprintf("请求参数序列化失败:%s %s", params, err.Error()))
	}
	toolRequest.Params.Arguments = params
	result, err := mcpClient.client.CallTool(context.Background(), toolRequest)
	if err != nil {
		return "", errors.New("工具调用失败" + err.Error())
	}
	return result.Content[0].(mcp.TextContent).Text, nil
}

type McpClient struct {
	name   string
	tools  []openai.Tool
	client *client.Client
}

func ClientMcp(cmd string, env []string, args []string, mcpType int) (*client.Client, error) {
	var err error
	//var resp *mcp.InitializeResult
	var c *client.Client
	c, err = createMcpClient(cmd, env, args, mcpType)
	if err != nil {
		return nil, err
	}
	_, err = c.Initialize(context.Background(), mcp.InitializeRequest{})
	return c, err
}

func createMcpClient(cmd string, env []string, args []string, mcpType int) (*client.Client, error) {
	if mcpType == 1 {
		return client.NewStdioMCPClient(cmd, env, args...)
	}
	c, err := client.NewSSEMCPClient(cmd)
	if err != nil {
		return nil, err
	}
	return c, c.Start(context.Background())
}

func Ping(cmd string, env []string, args []string, mcpType int) error {
	c, err := ClientMcp(cmd, env, args, mcpType)
	if err != nil {
		return err
	}
	defer c.Close()
	return c.Ping(context.Background())
}

func McpPing(c *client.Client) error {
	return c.Ping(context.Background())
}

func ConvertMCPToolsToOpenAI(mcps []model.Mcp) McpClientList {
	clients := make([]*McpClient, 0)
	for _, val := range mcps {
		c, err := ClientMcp(val.Cmd, strings.Split(val.Env, ","), strings.Split(val.Params, ","), val.Type)
		if err != nil {
			continue
		}
		if McpPing(c) != nil {
			continue
		}
		//如果too
		result, err := c.ListTools(context.Background(), mcp.ListToolsRequest{})
		if err != nil || len(result.Tools) == 0 {
			continue
		}
		tools := make([]openai.Tool, 0)
		for _, tool := range result.Tools {
			schema := tool.InputSchema
			tools = append(tools, openai.Tool{
				Type: openai.ToolTypeFunction,
				Function: &openai.FunctionDefinition{
					Description: tool.Description,
					Name:        val.Name + "-" + tool.Name,
					Parameters: map[string]interface{}{
						"type":       schema.Type,
						"properties": schema.Properties,
						"required":   schema.Required,
					},
				},
			})
		}
		clients = append(clients, &McpClient{
			name:   val.Name,
			tools:  tools,
			client: c,
		})
	}
	return clients
}
